#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
一键抠图工具图形界面
"""

import sys
import os
import threading
from pathlib import Path
from typing import List, Optional

try:
    import tkinter as tk
    from tkinter import ttk, filedialog, messagebox
except ImportError as e:
    print("❌ 缺少GUI依赖，请安装tkinter（通常Python自带）")
    sys.exit(1)

try:
    from .remover import BackgroundRemover
except ImportError:
    # 如果相对导入失败，尝试绝对导入
    try:
        from remover import BackgroundRemover
    except ImportError:
        # 如果仍然失败，尝试从父目录导入
        import sys
        from pathlib import Path
        current_dir = Path(__file__).parent
        sys.path.insert(0, str(current_dir))
        from remover import BackgroundRemover


class BackgroundRemoverGUI:
    """抠图工具图形界面"""
    
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("一键抠图工具 v2.0")
        self.root.geometry("700x600")  # 恢复原来的窗口大小
        self.root.resizable(True, True)
        self.root.minsize(600, 500)  # 恢复原来的最小窗口大小
        
        # 设置图标和样式
        self.setup_style()
        
        # 初始化变量
        self.selected_files = []
        self.output_directory = tk.StringVar()
        self.model_var = tk.StringVar(value="u2net")
        self.progress_var = tk.DoubleVar()
        self.status_var = tk.StringVar(value="就绪 - 请选择图片文件开始抠图")
        
        # 设置默认输出目录
        self.set_default_output_dir()
        
        # 创建界面
        self.create_widgets()
        
        # 居中显示窗口
        self.center_window()
    
    def setup_style(self):
        """设置界面样式"""
        style = ttk.Style()
        style.theme_use('clam')
        
        # 配置样式
        style.configure('Title.TLabel', font=('Arial', 16, 'bold'))
        style.configure('Heading.TLabel', font=('Arial', 10, 'bold'))
        style.configure('Status.TLabel', font=('Arial', 9))
    
    def set_default_output_dir(self):
        """设置默认输出目录"""
        desktop_dir = Path.home() / "Desktop" / "抠图输出"
        self.output_directory.set(str(desktop_dir))
    
    def center_window(self):
        """使窗口居中显示"""
        self.root.update_idletasks()
        width = self.root.winfo_width()
        height = self.root.winfo_height()
        x = (self.root.winfo_screenwidth() // 2) - (width // 2)
        y = (self.root.winfo_screenheight() // 2) - (height // 2)
        self.root.geometry(f'{width}x{height}+{x}+{y}')
    
    def create_widgets(self):
        """创建界面组件"""
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky="nsew")
        
        # 配置网格权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        
        # 标题
        title_label = ttk.Label(main_frame, text="AI智能抠图工具", style='Title.TLabel')
        title_label.grid(row=0, column=0, columnspan=3, pady=(0, 10))
        
        # 文件选择区域
        self.create_file_selection_area(main_frame)
        
        # 输出目录选择区域
        self.create_output_selection_area(main_frame)
        
        # 模型选择区域
        self.create_model_selection_area(main_frame)
        
        # 操作按钮区域
        self.create_action_buttons(main_frame)
        
        # 进度条
        self.create_progress_area(main_frame)
        
        # 状态栏
        self.create_status_bar(main_frame)
    
    def create_file_selection_area(self, parent):
        """创建文件选择区域"""
        row = 1
        
        # 标题
        ttk.Label(parent, text="选择要处理的图片:", style='Heading.TLabel').grid(
            row=row, column=0, columnspan=3, sticky="w", pady=(0, 5)
        )
        
        # 文件列表框架
        list_frame = ttk.Frame(parent)
        list_frame.grid(row=row+1, column=0, columnspan=3, sticky="ew", pady=(0, 10))
        list_frame.columnconfigure(0, weight=1)
        
        # 文件列表
        self.file_listbox = tk.Listbox(list_frame, height=8, selectmode=tk.EXTENDED)
        self.file_listbox.grid(row=0, column=0, sticky="ew")
        
        # 滚动条
        scrollbar = ttk.Scrollbar(list_frame, orient="vertical", command=self.file_listbox.yview)
        scrollbar.grid(row=0, column=1, sticky="ns")
        self.file_listbox.config(yscrollcommand=scrollbar.set)
        
        # 文件操作按钮
        button_frame = ttk.Frame(parent)
        button_frame.grid(row=row+2, column=0, columnspan=3, pady=(0, 10))
        
        ttk.Button(button_frame, text="选择图片文件", command=self.select_files).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="选择图片文件夹", command=self.select_folder).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="清空列表", command=self.clear_files).pack(side=tk.LEFT, padx=(0, 5))
        ttk.Button(button_frame, text="移除选中", command=self.remove_selected).pack(side=tk.LEFT)
    
    def create_output_selection_area(self, parent):
        """创建输出目录选择区域"""
        row = 4
        
        # 标题
        ttk.Label(parent, text="输出目录:", style='Heading.TLabel').grid(
            row=row, column=0, sticky="w", pady=(0, 5)
        )
        
        # 输出目录输入框和按钮
        output_frame = ttk.Frame(parent)
        output_frame.grid(row=row+1, column=0, columnspan=3, sticky="ew", pady=(0, 10))
        output_frame.columnconfigure(0, weight=1)
        
        ttk.Entry(output_frame, textvariable=self.output_directory, width=50).grid(
            row=0, column=0, sticky="ew", padx=(0, 5)
        )
        ttk.Button(output_frame, text="选择目录", command=self.select_output_dir).grid(
            row=0, column=1
        )
    
    def create_model_selection_area(self, parent):
        """创建AI模型选择区域"""
        row = 6
        
        # 标题
        ttk.Label(parent, text="AI模型选择:", style='Heading.TLabel').grid(
            row=row, column=0, sticky="w", pady=(0, 5)
        )
        
        # 模型选择
        model_frame = ttk.Frame(parent)
        model_frame.grid(row=row+1, column=0, columnspan=3, sticky="w", pady=(0, 10))
        
        models = [
            ("u2net", "U²Net (推荐) - 通用效果最佳"),
            ("u2net_human_seg", "U²Net-Human - 专门用于人像"),
            ("silueta", "Silueta - 高精度边缘检测")
        ]
        
        for i, (value, text) in enumerate(models):
            ttk.Radiobutton(
                model_frame, text=text, variable=self.model_var, value=value
            ).grid(row=i, column=0, sticky="w", pady=2)
    
    def create_action_buttons(self, parent):
        """创建操作按钮"""
        row = 10
        
        # 操作按钮标题
        ttk.Label(parent, text="操作控制:", style='Heading.TLabel').grid(
            row=row, column=0, columnspan=3, sticky="w", pady=(10, 5)
        )
        
        button_frame = ttk.Frame(parent)
        button_frame.grid(row=row+1, column=0, columnspan=3, pady=(0, 50))
        
        # 开始处理按钮 - 使用更明显的样式
        self.process_button = ttk.Button(
            button_frame, 
            text="开始抠图", 
            command=self.start_processing,
            width=15
        )
        self.process_button.pack(side=tk.LEFT, padx=(0, 10))
        
        # 停止按钮
        self.stop_button = ttk.Button(
            button_frame, 
            text="停止处理", 
            command=self.stop_processing, 
            state=tk.DISABLED,
            width=15
        )
        self.stop_button.pack(side=tk.LEFT, padx=(0, 10))
        
        # 打开输出目录按钮
        ttk.Button(
            button_frame, 
            text="打开输出目录", 
            command=self.open_output_dir,
            width=15
        ).pack(side=tk.LEFT)
    
    def create_progress_area(self, parent):
        """创建进度条区域"""
        row = 11
        
        # 进度条
        self.progress_bar = ttk.Progressbar(
            parent, mode='determinate', variable=self.progress_var
        )
        self.progress_bar.grid(row=row, column=0, columnspan=3, sticky="ew", pady=10)
    
    def create_status_bar(self, parent):
        """创建状态栏"""
        row = 12
        
        # 状态标签
        status_label = ttk.Label(parent, textvariable=self.status_var, style='Status.TLabel')
        status_label.grid(row=row, column=0, columnspan=3, sticky="w")
    
    def select_files(self):
        """选择图片文件"""
        filetypes = [
            ("图片文件", "*.jpg *.jpeg *.png *.bmp *.tiff *.webp"),
            ("JPEG文件", "*.jpg *.jpeg"),
            ("PNG文件", "*.png"),
            ("所有文件", "*.*")
        ]
        
        files = filedialog.askopenfilenames(
            title="选择图片文件",
            filetypes=filetypes
        )
        
        for file in files:
            if file not in self.selected_files:
                self.selected_files.append(file)
                self.file_listbox.insert(tk.END, Path(file).name)
        
        self.update_status(f"已选择 {len(self.selected_files)} 个文件")
    
    def select_folder(self):
        """选择图片文件夹"""
        folder = filedialog.askdirectory(title="选择包含图片的文件夹")
        
        if folder:
            folder_path = Path(folder)
            supported_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp'}
            
            new_files = []
            for file_path in folder_path.iterdir():
                if file_path.is_file() and file_path.suffix.lower() in supported_extensions:
                    file_str = str(file_path)
                    if file_str not in self.selected_files:
                        self.selected_files.append(file_str)
                        self.file_listbox.insert(tk.END, file_path.name)
                        new_files.append(file_path.name)
            
            if new_files:
                self.update_status(f"从文件夹添加了 {len(new_files)} 个图片文件")
            else:
                self.update_status("文件夹中没有找到新的图片文件")
    
    def clear_files(self):
        """清空文件列表"""
        self.selected_files.clear()
        self.file_listbox.delete(0, tk.END)
        self.update_status("文件列表已清空")
    
    def remove_selected(self):
        """移除选中的文件"""
        selected_indices = self.file_listbox.curselection()
        
        # 从后往前删除，避免索引变化
        for index in reversed(selected_indices):
            self.file_listbox.delete(index)
            del self.selected_files[index]
        
        self.update_status(f"已移除 {len(selected_indices)} 个文件")
    
    def select_output_dir(self):
        """选择输出目录"""
        directory = filedialog.askdirectory(title="选择输出目录")
        if directory:
            self.output_directory.set(directory)
            self.update_status(f"输出目录已设置为: {Path(directory).name}")
    
    def start_processing(self):
        """开始处理图片"""
        if not self.selected_files:
            messagebox.showwarning("提示", "请先选择要处理的图片文件！\n\n可以通过以下方式选择:\n• 点击'选择图片文件'选择单个或多个文件\n• 点击'选择图片文件夹'批量选择整个文件夹")
            return
        
        output_dir = self.output_directory.get()
        if not output_dir:
            messagebox.showwarning("提示", "请选择输出目录！")
            return
        
        # 显示详细的确认信息
        model_names = {
            "u2net": "U²Net (推荐) - 通用效果最佳",
            "u2net_human_seg": "U²Net-Human - 专门用于人像",
            "silueta": "Silueta - 高精度边缘检测"
        }
        
        result = messagebox.askyesno(
            "确认开始处理",
            f"处理设置确认:\n\n"
            f"📁 图片数量: {len(self.selected_files)} 个文件\n"
            f"📂 输出目录: {output_dir}\n"
            f"🤖 AI模型: {model_names.get(self.model_var.get(), self.model_var.get())}\n\n"
            f"⏱️ 首次使用需要下载AI模型，请耐心等待\n\n"
            "确定开始处理吗？"
        )
        
        if not result:
            return
        
        # 禁用处理按钮，启用停止按钮
        self.process_button.config(state=tk.DISABLED, text="处理中...")
        self.stop_button.config(state=tk.NORMAL)
        
        # 重置进度条
        self.progress_var.set(0)
        self.processing_stopped = False
        
        # 在新线程中处理
        self.processing_thread = threading.Thread(target=self.process_images, daemon=True)
        self.processing_thread.start()
    
    def stop_processing(self):
        """停止处理"""
        self.processing_stopped = True
        self.update_status("正在停止处理...")
    
    def process_images(self):
        """处理图片的工作线程"""
        try:
            output_dir = Path(self.output_directory.get())
            
            # 创建输出目录
            output_dir.mkdir(parents=True, exist_ok=True)
            
            # 初始化背景移除器
            self.update_status("正在加载AI模型...")
            remover = BackgroundRemover(self.model_var.get())
            
            success_count = 0
            total_count = len(self.selected_files)
            
            for i, file_path in enumerate(self.selected_files):
                if self.processing_stopped:
                    break
                
                input_path = None
                try:
                    input_path = Path(file_path)
                    output_path = output_dir / f"{input_path.stem}_no_bg.png"
                    
                    self.update_status(f"正在处理: {input_path.name} ({i+1}/{total_count})")
                    
                    # 处理图片
                    remover.remove_background(input_path, output_path)
                    success_count += 1
                    
                    # 更新进度条
                    progress = ((i + 1) / total_count) * 100
                    self.root.after(0, lambda p=progress: self.progress_var.set(p))
                    
                except Exception as e:
                    file_name = input_path.name if input_path else Path(file_path).name
                    print(f"处理失败 {file_name}: {e}")
            
            # 处理完成
            if self.processing_stopped:
                self.update_status("处理已停止")
            else:
                self.update_status(f"处理完成！成功: {success_count}/{total_count}")
                self.root.after(0, lambda: messagebox.showinfo(
                    "处理完成",
                    f"图片处理完成！\n\n"
                    f"成功处理: {success_count} 个文件\n"
                    f"输出目录: {output_dir}\n\n"
                    f"点击「打开输出目录」查看结果"
                ))
        
        except Exception as e:
            error_msg = f"处理过程中发生错误: {e}"
            self.update_status(error_msg)
            self.root.after(0, lambda: messagebox.showerror("错误", error_msg))
        
        finally:
            # 恢复按钮状态
            self.root.after(0, self.reset_buttons)
    
    def reset_buttons(self):
        """重置按钮状态"""
        self.process_button.config(state=tk.NORMAL, text="开始抠图")
        self.stop_button.config(state=tk.DISABLED)
    
    def open_output_dir(self):
        """打开输出目录"""
        output_dir = self.output_directory.get()
        if output_dir and Path(output_dir).exists():
            os.startfile(output_dir)  # Windows
        else:
            messagebox.showwarning("警告", "输出目录不存在")
    
    def update_status(self, message):
        """更新状态信息"""
        def update():
            self.status_var.set(message)
            print(f"[状态] {message}")
        
        if threading.current_thread() != threading.main_thread():
            self.root.after(0, update)
        else:
            update()
    
    def run(self):
        """运行GUI"""
        self.root.mainloop()


def show_gui_progress_bar(is_exe_mode=True, duration=1.5, steps=15):
    """显示GUI初始化进度条"""
    import time
    
    prefix = "     " if is_exe_mode else "  "
    
    for i in range(steps + 1):
        progress = i / steps
        filled_length = int(progress * 20)
        bar = '█' * filled_length + '-' * (20 - filled_length)
        percent = int(progress * 100)
        
        print(f"\r{prefix}[█{bar}█] {percent:3d}%", end='', flush=True)
        
        if i < steps:
            time.sleep(duration / steps)
    
    print()  # 换行

def main():
    """启动图形界面"""
    # 检测是否运行在exe模式下
    is_exe_mode = getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS')
    
    try:
        if is_exe_mode:
            print("   🎨 初始化GUI组件...")
            show_gui_progress_bar(is_exe_mode, duration=2.0)
        else:
            print("🎨 初始化GUI组件...")
            show_gui_progress_bar(is_exe_mode, duration=1.0)
            
        app = BackgroundRemoverGUI()
        
        if is_exe_mode:
            print("   ✅ GUI组件初始化完成")
            print("   🚀 窗口即将显示...")
            show_gui_progress_bar(is_exe_mode, duration=1.0)
        else:
            print("✅ GUI组件初始化完成")
            print("🚀 窗口即将显示...")
            show_gui_progress_bar(is_exe_mode, duration=0.5)
            
        app.run()
    except Exception as e:
        print(f"GUI启动失败: {e}")
        input("按Enter键退出...")


if __name__ == "__main__":
    main()